/**
 * @file all_gather_matmul_custom.cpp
 *
 * Copyright (C) 2024. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */

#include "../op_kernel/all_gather_matmul_custom_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include "tiling/tiling_api.h"
#include "register/tilingdata_base.h"
#include "tiling/hccl/hccl_tiling.h"

#define INFO_LOG(fmt, args...) fprintf(stdout, "[INFO]  " fmt "\n", ##args)
#define WARN_LOG(fmt, args...) fprintf(stdout, "[WARN]  " fmt "\n", ##args)
#define ERROR_LOG(fmt, args...) fprintf(stderr, "[ERROR]  " fmt "\n", ##args)

// tiling
constexpr int32_t TILE_M = 448;
constexpr uint32_t HCCL_CMD_ALLGATHER = 6;
constexpr uint32_t FP16_BF16_SIZE = 2;
constexpr uint32_t CUSTOM_TILING_KEY = 100; // full mesh + no nd2nz + no cast bias
constexpr int32_t L1_BUFFER_SIZE = 512 * 1024;

namespace {
enum class HcclDataType {
    HCCL_DATA_TYPE_INT8 = 0,   /* *< int8 */
    HCCL_DATA_TYPE_INT16 = 1,  /* *< int16 */
    HCCL_DATA_TYPE_INT32 = 2,  /* *< int32 */
    HCCL_DATA_TYPE_FP16 = 3,   /* *< fp16 */
    HCCL_DATA_TYPE_FP32 = 4,   /* *< fp32 */
    HCCL_DATA_TYPE_INT64 = 5,  /* *< int64 */
    HCCL_DATA_TYPE_UINT64 = 6, /* *< uint64 */
    HCCL_DATA_TYPE_UINT8 = 7,  /* *< uint8 */
    HCCL_DATA_TYPE_UINT16 = 8, /* *< uint16 */
    HCCL_DATA_TYPE_UINT32 = 9, /* *< uint32 */
    HCCL_DATA_TYPE_FP64 = 10,  /* *< fp64 */
    HCCL_DATA_TYPE_BFP16 = 11, /* *< bfp16 */
    HCCL_DATA_TYPE_RESERVED    /* *< reserved */
};

const std::set<ge::DataType> SUPPORT_DTYPE = { ge::DT_FLOAT16, ge::DT_BF16 };
const std::map<ge::DataType, matmul_tiling::DataType> DTYPE_MAP = {
    {ge::DT_FLOAT16, matmul_tiling::DataType::DT_FLOAT16},
    {ge::DT_BF16, matmul_tiling::DataType::DT_BFLOAT16}
};

const std::map<ge::DataType, HcclDataType> HCCL_DATA_TYPE_MAP = {
    {ge::DataType::DT_FLOAT16, HcclDataType::HCCL_DATA_TYPE_FP16},
    {ge::DataType::DT_BF16, HcclDataType::HCCL_DATA_TYPE_BFP16}
};
}

static ge::graphStatus AllGatherMatmulCustomTilingFunc(gert::TilingContext *context)
{
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    auto aicCoreNum = ascendcPlatform.GetCoreNumAic();

    // get attrs
    uint32_t index = 0U;
    const char *group = context->GetAttrs()->GetAttrPointer<char>(index++);
    bool isTransA = *(context->GetAttrs()->GetAttrPointer<bool>(index++));
    bool isTransB = *(context->GetAttrs()->GetAttrPointer<bool>(index++));
    int gatherIndex = *(context->GetAttrs()->GetAttrPointer<int>(index++));
    int commTurn = *(context->GetAttrs()->GetAttrPointer<int>(index++));
    int rankSize = *(context->GetAttrs()->GetAttrPointer<int>(index++));
    bool isGatherOut = *(context->GetAttrs()->GetAttrPointer<int>(index++));
    INFO_LOG("Group %s, isTransA is %d, isTransB is %d, gatherIndex is %d, commTurn is %d, rankSize %d, isGatherOut %d",
        group, isTransA, isTransB, gatherIndex, commTurn, rankSize, isGatherOut);

    // get shape  [[4096/8,5120], [5120,640]] fp16
    uint64_t rankM = context->GetInputShape(0)->GetStorageShape().GetDim(0);
    uint64_t rankK = context->GetInputShape(0)->GetStorageShape().GetDim(1);
    uint64_t rankN = isTransB ?
        context->GetInputShape(1)->GetStorageShape().GetDim(0) : context->GetInputShape(1)->GetStorageShape().GetDim(1);
    INFO_LOG("RankM %lu, rankK %lu, rankN %lu", rankM, rankK, rankN);

    // get dtype
    auto aType = context->GetInputTensor(0)->GetDataType();
    auto bType = context->GetInputTensor(1)->GetDataType();
    auto cType = aType;
    if (SUPPORT_DTYPE.find(aType) == SUPPORT_DTYPE.cend() || SUPPORT_DTYPE.find(bType) == SUPPORT_DTYPE.cend()) {
        ERROR_LOG("Dtype of a %d or b %d is unsupported", static_cast<int>(aType), static_cast<int>(bType));
        return ge::GRAPH_FAILED;
    }

    // set block dim & tiling key
    context->SetBlockDim(ascendcPlatform.GetCoreNumAic());
    context->SetTilingKey(CUSTOM_TILING_KEY);

    // set work space size
    size_t nd2nzLen = (rankN + 16) * (rankK + 16) * FP16_BF16_SIZE;
    size_t systemWorkspaceSize = static_cast<size_t>(ascendcPlatform.GetLibApiWorkSpaceSize());
    size_t workspaceSize = nd2nzLen + 0 + 0 + 0 + systemWorkspaceSize; // nd2nzLen + gmcFloat + gatherLen + biasLen + 16M
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = workspaceSize;

    uint64_t tileNum = rankM / TILE_M;
    uint64_t tailNum = (rankM % TILE_M == 0) ? 0 : 1;
    uint64_t tailM = rankM % TILE_M;
    INFO_LOG("tileNum %lu, tailNum %lu, tailM %lu", tileNum, tailNum, tailM);

    AllGatherMatmulCustomTilingData *tiling = context->GetTilingData<AllGatherMatmulCustomTilingData>();

    tiling->param.rankDim = rankSize;
    tiling->param.tileNum = tileNum;
    tiling->param.tailM = tailM;
    tiling->param.tailNum = tailNum;
    tiling->param.rankM = rankM;
    tiling->param.rankN = rankN;
    tiling->param.rankK = rankK;
    tiling->param.gatherIndex = gatherIndex;
    tiling->param.isTransposeA = isTransA ? 1 : 0;
    tiling->param.isTransposeB = isTransB ? 1 : 0;
    tiling->param.storageGather = isGatherOut;
    tiling->param.cToFloatLen = 0; // 通信结果输出到workspace或gatherOut 不需要gmcFloat
    tiling->param.nd2NzWorkLen = nd2nzLen;
    tiling->param.gatherLen = 0; // gatherOut=true 不输出到workspace
    tiling->param.dataType = static_cast<uint8_t>(HCCL_DATA_TYPE_MAP.at(aType));
    // matmul tiling func
    auto matmulTilingFunc = [&](int64_t m, int64_t n, int64_t k, TCubeTiling &cubeTiling) -> bool {
        matmul_tiling::MultiCoreMatmulTiling mmTiling;
        mmTiling.SetAType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, DTYPE_MAP.at(aType), isTransA);
        mmTiling.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, DTYPE_MAP.at(bType), isTransB);
        mmTiling.SetCType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, DTYPE_MAP.at(cType));
        mmTiling.SetBias(false);
        mmTiling.SetDim(aicCoreNum);
        mmTiling.SetShape(m, n, k);
        mmTiling.SetOrgShape(m, n, k);
        mmTiling.SetBufferSpace(L1_BUFFER_SIZE);
        if (mmTiling.GetTiling(cubeTiling) != 0) {
            return false;
        }
        return true;
    };
    // matmul local tiling
    if (!matmulTilingFunc(rankM, rankN, rankK, tiling->localTiling)) {
        ERROR_LOG("Get local matmul tiling failed");
        return ge::GRAPH_FAILED;
    }
    // matmul tile tiling
    if (!matmulTilingFunc(TILE_M, rankN, rankK, tiling->tileTiling)) {
        ERROR_LOG("Get tile matmul tiling failed");
        return ge::GRAPH_FAILED;
    }
    // matmul tail tiling
    if (!matmulTilingFunc(rankM % TILE_M, rankN, rankK, tiling->tailTiling)) {
        ERROR_LOG("Get tail matmul tiling failed");
        return ge::GRAPH_FAILED;
    }

    uint32_t opType = 6;
    std::string algConfig = "AllGather=level0:doublering";
    uint32_t reduceType = 0;
    AscendC::Mc2CcTilingConfig mc2CcTilingConfig(group, opType, algConfig, reduceType);
    mc2CcTilingConfig.GetTiling(tiling->mc2InitTiling);
    mc2CcTilingConfig.SetSkipLocalRankCopy(0);
    mc2CcTilingConfig.GetTiling(tiling->mc2CcTiling);
    
    return ge::GRAPH_SUCCESS;
}

namespace ops {
class AllGatherMatmulCustom : public OpDef {
public:
    explicit AllGatherMatmulCustom(const char *name) : OpDef(name)
    {
        this->Input("x1")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("x2")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND})
            .IgnoreContiguous();
        this->Input("bias")
            .ParamType(OPTIONAL)
            .DataType({ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("gather_out")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

        this->Attr("group").AttrType(REQUIRED).String();
        this->Attr("isTransA").AttrType(OPTIONAL).Bool(false);
        this->Attr("isTransB").AttrType(OPTIONAL).Bool(false);
        this->Attr("gatherIndex").AttrType(OPTIONAL).Int(0);
        this->Attr("commTurn").AttrType(OPTIONAL).Int(0);
        this->Attr("rank_size").AttrType(OPTIONAL).Int(0);
        this->Attr("is_gather_out").AttrType(OPTIONAL).Bool(true);

        OpAICoreConfig aicore_config;
        aicore_config.DynamicCompileStaticFlag(true)
            .DynamicFormatFlag(true)
            .DynamicRankSupportFlag(true)
            .DynamicShapeSupportFlag(true)
            .NeedCheckSupportFlag(false)
            .PrecisionReduceFlag(true)
            .ExtendCfgInfo("aclnnSupport.value", "support_aclnn")
            .ExtendCfgInfo("jitCompile.flag", "static_false")
            .ExtendCfgInfo("multiKernelSupportDynamicGraph.value", "multi_kernel");
        this->AICore().SetTiling(AllGatherMatmulCustomTilingFunc);
        this->AICore().AddConfig("ascend910b", aicore_config);
        this->MC2().HcclGroup("group");
    }
};

OP_ADD(AllGatherMatmulCustom);
}
